/*
 * @file pdmat.h
 *
 * The class to represent positive semi-definite matrices
 *
 * @author Dahua Lin
 */

#ifndef _MSC_VER
#pragma once
#endif

#ifndef GMLPP_PDMAT_H_
#define GMLPP_PDMAT_H_

#include <gmlpp/base/arma.h>

namespace gml
{
	/********************************************
	 *
	 *  base class
	 *
	 ********************************************/

	template<class Derived, typename T>
	class PDMatBase
	{
	public:
		typedef T value_type;

		const Derived& derived() const
		{
			return static_cast<const Derived&>(*this);
		}

	public:

		uword dim() const
		{
			return derived().dim();
		}

		T log_det() const
		{
			return derived().log_det();
		}

		T trace() const
		{
			return derived().trace();
		}

		void get_mat(Mat<T>& R) const
		{
			derived().get_mat(R);
		}

		Mat<T> get_mat() const
		{
			return derived().get_mat();
		}

		// matrix product

		void mtimes(const Mat<T>& X, Mat<T>& Y) const
		{
			derived().mtimes(X, Y);
		}

		Mat<T> operator * (const Mat<T>& X) const
		{
			return derived().operator *(X);
		}

		Col<T> operator * (const Col<T>& x) const
		{
			return derived().operator *(x);
		}

		// solve

		void solve(const Mat<T>& X, Mat<T>& Y) const
		{
			derived().solve(X, Y);
		}

		Mat<T> solve(const Mat<T>& X) const
		{
			return derived().solve(X);
		}

		Col<T> solve(const Col<T>& x) const
		{
			return derived().solve(x);
		}

		// quad

		T quad(const Col<T>& x) const
		{
			return derived().quad(x);
		}

		T quad(const Col<T>& x, const Col<T>& y) const
		{
			return derived().quad(x, y);
		}

		void colwise_quad(const Mat<T>& X, const Mat<T>& Y, Row<T>& r) const
		{
			derived().colwise_quad(X, Y, r);
		}

		Row<T> colwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return derived().colwise_quad(X, Y);
		}

		void pairwise_quad(const Mat<T>& X, const Mat<T>& Y, Mat<T>& R) const
		{
			derived().pairwise_quad(X, Y, R);
		}

		Mat<T> pairwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return derived().pairwise_quad(X, Y);
		}

		Derived inv() const
		{
			return derived().inv();
		}

	}; // end class PDMatBase



	/********************************************
	 *
	 *  isotropic diagonal matrix
	 *
	 ********************************************/

	template<typename T>
	class PIsoMat : public PDMatBase<PIsoMat<T>, T>
	{
	public:
		PIsoMat(uword d, T dv)
		: m_d(d), m_dv(dv), m_invdv(T(1) / dv)
		{
		}

		T diagval() const
		{
			return m_dv;
		}

	public:
		uword dim() const
		{
			return m_d;
		}

		T log_det() const
		{
			return log(m_dv) * m_d;
		}

		T trace() const
		{
			return m_dv * m_d;
		}

		Mat<T> get_mat() const
		{
			Col<T> dvec(m_d);
			dvec.fill(m_dv);
			return arma::diagmat(dvec);
		}

		void get_mat(Mat<T>& R) const
		{
			uword d = m_d;
			R.zeros(d, d);
			for (uword i = 0; i < d; ++i) R.at(i, i) = m_dv;
		}

		// matrix product

		void mtimes(const Mat<T>& X, Mat<T>& Y) const
		{
			Y = X * m_dv;
		}

		Mat<T> operator * (const Mat<T>& X) const
		{
			return X * m_dv;
		}

		Col<T> operator * (const Col<T>& x) const
		{
			return x * m_dv;
		}

		// solve

		void solve(const Mat<T>& X, Mat<T>& Y) const
		{
			Y = X * m_invdv;
		}

		Mat<T> solve(const Mat<T>& X) const
		{
			return X * m_invdv;
		}

		Col<T> solve(const Col<T>& x) const
		{
			return x * m_invdv;
		}

		// quad

		T quad(const Col<T>& x) const
		{
			return arma::dot(x, x) * m_dv;
		}

		T quad(const Col<T>& x, const Col<T>& y) const
		{
			return arma::dot(x, y) * m_dv;
		}

		void colwise_quad(const Mat<T>& X, const Mat<T>& Y, Row<T>& r) const
		{
			r = sum(X % Y, 0) * m_dv;
		}

		Row<T> colwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return sum(X % Y, 0) * m_dv;
		}

		void pairwise_quad(const Mat<T>& X, const Mat<T>& Y, Mat<T>& R) const
		{
			R =  (X.t() * Y) * m_dv;
		}

		Mat<T> pairwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return (X.t() * Y) * m_dv;
		}

		PIsoMat inv() const
		{
			return PIsoMat(m_d, math::rcp(m_dv));
		}

	private:
		uword m_d;
		T m_dv;
		T m_invdv;

	}; // end class PIsoMat




	/********************************************
	 *
	 *  diagonal matrix
	 *
	 ********************************************/

	template<typename T>
	class PDiagMat : public PDMatBase<PDiagMat<T>, T>
	{
	public:
		explicit PDiagMat(uword n)
		: m_dv(n)
		{
		}

		explicit PDiagMat(const Col<T>& dv)
		: m_dv(dv)
		{
		}

		const Col<T>& diagvals() const
		{
			return m_dv;
		}

		Col<T>& diagvals()
		{
			return m_dv;
		}

	public:
		uword dim() const
		{
			return m_dv.n_elem;
		}

		T log_det() const
		{
			Col<T> log_dv = arma::log(m_dv);
			return arma::sum(log_dv);
		}

		T trace() const
		{
			return arma::sum(m_dv);
		}

		Mat<T> get_mat() const
		{
			return arma::diagmat(m_dv);
		}

		void get_mat(Mat<T>& R) const
		{
			uword d = m_dv.n_elem;
			R.zeros(d, d);
			for (uword i = 0; i < d; ++i) R.at(i, i) = m_dv[i];
		}

		// matrix product

		void mtimes(const Mat<T>& X, Mat<T>& Y) const
		{
			if (X.n_cols == 1)
			{
				Y = X % m_dv;
			}
			else
			{
				bsxc_times(X, m_dv, Y);
			}
		}

		Mat<T> operator * (const Mat<T>& X) const
		{
			return bsxc_times(X, m_dv);
		}

		Col<T> operator * (const Col<T>& x) const
		{
			return x % m_dv;
		}

		// solve

		void solve(const Mat<T>& X, Mat<T>& Y) const
		{
			if (X.n_cols == 1)
			{
				Y = X % inv_dv();
			}
			else
			{
				bsxc_times(X, inv_dv(), Y);
			}
		}

		Mat<T> solve(const Mat<T>& X) const
		{
			return bsxc_times(X, inv_dv());
		}

		Col<T> solve(const Col<T>& x) const
		{
			return x % inv_dv();
		}

		// quad

		T quad(const Col<T>& x) const
		{
			return arma::dot(arma::square(x), m_dv);
		}

		T quad(const Col<T>& x, const Col<T>& y) const
		{
			return arma::dot((x % y), m_dv);
		}

		void colwise_quad(const Mat<T>& X, const Mat<T>& Y, Row<T>& r) const
		{
			r = m_dv.t() * (X % Y);
		}

		Row<T> colwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return m_dv.t() * (X % Y);
		}

		void pairwise_quad(const Mat<T>& X, const Mat<T>& Y, Mat<T>& R) const
		{
			R = X.t() * bsxc_times(Y, m_dv);
		}

		Mat<T> pairwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return X.t() * bsxc_times(Y, m_dv);
		}

		PDiagMat inv() const
		{
			uword d = dim();

			PDiagMat R(d);
			const T *s = m_dv.memptr();
			T *r = R.diagvals().memptr();

			for (uword i = 0; i < d; ++i) r[i] = math::rcp(s[i]);

			return R;
		}

	private:
		GML_ENSURE_INLINE Col<T> inv_dv() const
		{
			return T(1) / m_dv;
		}

	private:
		Col<T> m_dv;

	}; // end class PDiagMat




	/********************************************
	 *
	 *  full matrix
	 *
	 ********************************************/

	template<typename T>
	class PFullMat : public PDMatBase<PFullMat<T>, T>
	{
	public:
		explicit PFullMat(uword n)
		: m_mat(n, n)
		{
		}

		PFullMat(const Mat<T>& M)
		: m_mat(M)
		{
			check_arg(M.n_rows == M.n_cols);
		}

		const Mat<T>& mat() const
		{
			return m_mat;
		}

		Mat<T>& mat()
		{
			return m_mat;
		}

	public:
		uword dim() const
		{
			return m_mat.n_rows;
		}

		T log_det() const
		{
			Mat<T> R = arma::chol(m_mat);
			return arma::sum(arma::log(arma::abs(R.diag()))) * 2;
		}

		T trace() const
		{
			return arma::trace(m_mat);
		}

		void get_mat(Mat<T>& R) const
		{
			R = m_mat;
		}

		Mat<T> get_mat() const
		{
			return m_mat;
		}

		// matrix product

		void mtimes(const Mat<T>& X, Mat<T>& Y) const
		{
			Y = m_mat * X;
		}

		Mat<T> operator * (const Mat<T>& X) const
		{
			return m_mat * X;
		}

		Col<T> operator * (const Col<T>& x) const
		{
			return m_mat * x;
		}

		// solve

		void solve(const Mat<T>& X, Mat<T>& Y) const
		{
			arma::solve(Y, m_mat, X);
		}

		Mat<T> solve(const Mat<T>& X) const
		{
			return arma::solve(m_mat, X);
		}

		Col<T> solve(const Col<T>& x) const
		{
			return arma::solve(m_mat, x);
		}

		T quad(const Col<T>& x) const
		{
			return arma::as_scalar(x.t() * m_mat * x);
		}

		T quad(const Col<T>& x, const Col<T>& y) const
		{
			return arma::as_scalar(x.t() * m_mat * y);
		}

		void colwise_quad(const Mat<T>& X, const Mat<T>& Y, Row<T>& r) const
		{
			r = arma::sum(X % (m_mat * Y), 0);
		}

		Row<T> colwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return arma::sum(X % (m_mat * Y), 0);
		}

		void pairwise_quad(const Mat<T>& X, const Mat<T>& Y, Mat<T>& R) const
		{
			R = X.t() * m_mat * Y;
		}

		Mat<T> pairwise_quad(const Mat<T>& X, const Mat<T>& Y) const
		{
			return X.t() * m_mat * Y;
		}

		PFullMat inv() const
		{
			PFullMat R(dim());
			arma::inv(R.m_mat, m_mat);
			return R;
		}

	private:
		Mat<T> m_mat;

	}; // end class PFullMat



}

#endif /* PDMAT_H_ */
